import open3d as o3d
import numpy as np
import copy

# 文件路径（交换源和目标点云）
source_path = "dummy2.ply"  # 源点云文件（原目标点云）
target_path = "dummy1.ply"  # 目标点云文件（原源点云）

# 1. 加载点云
source = o3d.io.read_point_cloud(source_path)
target = o3d.io.read_point_cloud(target_path)

# 2. 估计法线（如果尚未估计）
def estimate_normals(pcd, radius=0.1, max_nn=30):
    pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=max_nn))
    pcd.normalize_normals()

# 检查并估计法线
if not source.has_normals():
    estimate_normals(source)

if not target.has_normals():
    estimate_normals(target)

# 3. 确保法线方向一致
def orient_normals_consistently(pcd):
    # 使用PCA来确定法线的大致方向
    points = np.asarray(pcd.points)
    if len(points) == 0:
        return
    centroid = np.mean(points, axis=0)
    cov = np.cov(points.T)
    eigenvalues, eigenvectors = np.linalg.eigh(cov)
    principal_direction = eigenvectors[:, -1]  # 主方向
    # 计算法线与主方向的点积
    normals = np.asarray(pcd.normals)
    dot_products = np.dot(normals, principal_direction)
    # 如果大多数法线与主方向的点积为负，则翻转所有法线
    if np.mean(dot_products) < 0:
        pcd.normals = o3d.utility.Vector3dVector(-normals)

orient_normals_consistently(source)
orient_normals_consistently(target)

# 4. 检查是否需要镜像
def is_mirrored(source_pcd, target_pcd, threshold=0.5):
    # 计算两个点云的法线方向是否相反
    source_normal = np.mean(np.asarray(source_pcd.normals), axis=0)
    target_normal = np.mean(np.asarray(target_pcd.normals), axis=0)
    dot_product = np.dot(source_normal, target_normal)
    return dot_product < threshold

mirrored = is_mirrored(source, target)
if mirrored:
    print("检测到点云可能存在镜像关系，进行镜像转换...")
    # 镜像转换：沿X轴翻转（根据具体情况选择轴）
    source_mirrored = copy.deepcopy(source)
    #source_mirrored.scale((1, 1, -1), center=(0, 0, 0))

    source_mirrored.scale(-1, center=(0,0,0))
    source = source_mirrored
    print("镜像转换完成。")

# 5. 初步对齐（使用点云中心对齐）
source_center = source.get_center()
target_center = target.get_center()
translation = target_center - source_center
source.translate(translation)
print(f"初步对齐完成，平移向量: {translation}")

# 6. 特征匹配进行初始对齐（可选，但推荐）
def preprocess_point_cloud(pcd, voxel_size):
    pcd_down = pcd.voxel_down_sample(voxel_size)
    pcd_down.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 2, max_nn=30))
    pcd_down, fpfh = compute_fpfh(pcd_down, voxel_size)
    return pcd_down, fpfh

def compute_fpfh(pcd_down, voxel_size):
    fpfh = o3d.pipelines.registration.compute_fpfh_feature(
        pcd_down,
        o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 5, max_nn=100))
    return pcd_down, fpfh

voxel_size = 0.05  # 根据点云大小调整
source_down, source_fpfh = preprocess_point_cloud(source, voxel_size)
target_down, target_fpfh = preprocess_point_cloud(target, voxel_size)

# 使用 RANSAC 进行全局配准
print("正在执行全局配准 (RANSAC)...")
result_ransac = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
    source_down, target_down, source_fpfh, target_fpfh, mutual_filter=True,
    max_correspondence_distance=voxel_size * 1.5,
    estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(False),
    ransac_n=4,
    checkers=[ 
        o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
        o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(voxel_size * 1.5)
    ],
    criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(4000000, 500)
)

print("RANSAC 全局配准完成。")
print("变换矩阵 (RANSAC):")
print(result_ransac.transformation)

# 将 RANSAC 变换应用到源点云
source.transform(result_ransac.transformation)

# 7. 使用 ICP 算法进行精细配准
threshold = 0.02  # 根据点云大小调整
print("正在执行 ICP 精细配准...")

reg_icp = o3d.pipelines.registration.registration_icp(
    source, target, threshold, result_ransac.transformation,
    o3d.pipelines.registration.TransformationEstimationPointToPoint()
)

print("ICP 精细配准完成。")
print("最终变换矩阵:")
print(reg_icp.transformation)

# 将 ICP 变换应用到源点云
source.transform(reg_icp.transformation)

# 8. 手动调整旋转（如果旋转部分仍然未完全对齐）
rotation_matrix = result_ransac.transformation[:3, :3]  # 获取旋转部分
translation_vector = result_ransac.transformation[:3, 3]  # 获取平移部分

# 手动旋转源点云
source.rotate(rotation_matrix, center=(0, 0, 0))
# 手动平移源点云
source.translate(translation_vector)

# 9. 可视化对齐结果（保留原始颜色）
print("正在可视化对齐结果...")
o3d.visualization.draw_geometries([source, target],
                                  zoom=0.455,
                                  front=[0.0, 0.0, -1.0],
                                  lookat=[0, 0, 0],
                                  up=[0, -1, 0],
                                  window_name="对齐结果")

# 10. 保存对齐后的点云（保留原始颜色）
aligned_source_path = "aligne11.ply"  # 根据交换后的源点云命名
o3d.io.write_point_cloud(aligned_source_path, source)
print(f"对齐后的源点云已保存到 {aligned_source_path}")
